
import torch
import numpy as np
torch.manual_seed(1)

def get_random_problems(batch_size, problem_size):
    problems = torch.rand(size=(batch_size, problem_size, 2))
    # problems.shape: (batch, problem, 2)
    return problems


def augment_xy_data_by_8_fold(problems):
    # problems.shape: (batch, problem, 2)

    x = problems[:, :, [0]]
    y = problems[:, :, [1]]
    # x,y shape: (batch, problem, 1)

    dat1 = torch.cat((x, y), dim=2)
    dat2 = torch.cat((1 - x, y), dim=2)
    dat3 = torch.cat((x, 1 - y), dim=2)
    dat4 = torch.cat((1 - x, 1 - y), dim=2)
    dat5 = torch.cat((y, x), dim=2)
    dat6 = torch.cat((1 - y, x), dim=2)
    dat7 = torch.cat((y, 1 - x), dim=2)
    dat8 = torch.cat((1 - y, 1 - x), dim=2)

    aug_problems = torch.cat((dat1, dat2, dat3, dat4, dat5, dat6, dat7, dat8), dim=0)
    # shape: (8*batch, problem, 2)

    return aug_problems


def get_local_problem(local_path):
    data = np.load(local_path,allow_pickle=True).item()
    min_x = data['min_x']
    max_x = data['max_x']
    data = list(data['data'])
    current_len = len(data)
    if current_len < 100:
        for i in range(100-current_len):
            data.append(data[0])
    data = np.array(data)
    print(data.shape)
    return torch.tensor(data, dtype=torch.float).unsqueeze(0), len(data), min_x, max_x